import torch
from diffusers import FluxPipeline

import os
current_dir = os.getcwd()
print(current_dir)
import sys
sys.path.insert(0,current_dir)
from MODEL_CKP import FLUX


pipe = FluxPipeline.from_pretrained(FLUX, torch_dtype=torch.bfloat16)
pipe.enable_model_cpu_offload() #save some VRAM by offloading the model to CPU. Remove this if you have enough GPU power
'''
麻花编织 twisted cable knit

蜂巢纹 honeycomb knit 

镂空花纹 lace knit / eyelet patterns 

鱼骨纹 herringbone knit 

波浪状针法 wave stitch 

螺旋交错的线迹 spiral or braided stitches 
'''
knit_styles = [
    'twisted cable knit',
    'honeycomb knit',
    'lace knit',
    'herringbone knit',
    'wave stitch',
    'spiral or braided stitches'
]
for knit_style in knit_styles:
    prompt = f"A knitted sweater featuring {knit_style} patterns across the front, soft wool texture, isolated on a plain white background."
    image = pipe(
        prompt,
        height=1024,
        width=1024,
        guidance_scale=3.5,
        num_inference_steps=50,
        max_sequence_length=512,
        generator=torch.Generator("cpu").manual_seed(0)
    ).images[0]
    image.save("flux-dev.png")